"""
AI服务，调用大模型API将自然语言转换为SQL
"""
from dashscope import Generation
import requests
import json
import config
from models import DatabaseSchema
from knowledge_base import KnowledgeBase

# 初始化知识库
knowledge_base = KnowledgeBase()

def detect_query_intent(question):
    """
    判断用户问题是数据库查询还是普通对话
    
    Args:
        question (str): 用户问题
    
    Returns:
        str: "database_query" 或 "general_conversation"
    """
    # 数据库相关关键词
    db_keywords = [
        'bank', 'finance', 'financial', 'account', 'loan', 'credit', 'debit', 
        'transaction', 'payment', 'deposit', 'withdraw', 'balance', 'interest',
        'mortgage', 'investment', 'fund', 'asset', 'liability', 'capital',
        '查询', '数据', '表', '记录', '统计', '计算', '多少', '哪些', 
        '数据库', 'select', 'where', 'from', '查找', '检索', 
        '导出', '分析', '列出', '显示', '风险', '客户', 
        '大于', '小于', '等于', '排序', '平均', '最大', '最小', '总和',
        '查看', '最近', '条件', '比较'
    ]
    
    # 计算问题中包含的数据库关键词数量
    keyword_count = sum(1 for keyword in db_keywords if keyword.lower() in question.lower())
    
    # 普通对话关键词
    general_keywords = [
        '你好', '介绍', '自己', '帮助', '你是谁', '你能做什么', 
        '功能', '使用', '说明', '解释', '指南', '教程', '版本',
        '谢谢', '感谢', '再见', '问候', '闲聊', 'hello', 'hi'
    ]
    
    # 检查是否有明确的普通对话关键词
    for keyword in general_keywords:
        if keyword.lower() in question.lower():
            return "general_conversation"
    
    # 如果包含数据库关键词或问题较长(可能是复杂查询)，判断为数据库查询
    if keyword_count > 0 or len(question) > 30:
        return "database_query"
        
    # 默认为普通对话
    return "general_conversation"

def generate_general_response(question):
    """
    生成普通对话的回答，无需SQL查询
    
    Args:
        question (str): 用户问题
    
    Returns:
        str: AI回答
    """
    # 提示词
    prompt = f"""
    {config.PROMPT_TEMPLATE_IDENTITY}
    现在，用户向你提出了一个普通问题(非数据库查询)：
    "{question}"
    
    请直接回答这个问题，注意：
    1. 不要生成SQL查询
    2. 如果被问及你的功能，请解释
    3. 保持友好专业的语气
    """
    # 使用选定的AI模型生成回复
    if config.ACTIVE_AI_MODEL == "deepseek":
        return call_deepseek_api(prompt) or "我是AI数据库问答助手，可以帮您查询数据，也可以回答一般问题。"
    else:
        return call_tongyi_api(prompt) or "我是AI数据库问答助手，可以帮您查询数据，也可以回答一般问题。"


def generate_prompt(question, db_schema, options):
    """生成增强的提示词"""
    # 准备数据库模式信息
    tables_info = ""
    for table in db_schema.tables:
        columns_info = []
        for col in table.columns:
            if col.enum_values:
                # 对于枚举类型，明确列出所有可能的值
                enum_values_str = ", ".join(f"'{v}'" for v in col.enum_values.keys())
                columns_info.append(f"{col.name} ({col.type}): {col.description} [可选值: {enum_values_str}]")
            else:
                columns_info.append(f"{col.name} ({col.type}): {col.description}")
        
        tables_info += f"Table: {table.name}\nColumns: {', '.join(columns_info)}\n\n"
    
    constraints = []
    if not options.get("allow_joins", False) and len(db_schema.tables) > 1:
        constraints.append("不要使用JOIN操作。")
    if not options.get("complex_aggregations", False):
        constraints.append("避免使用复杂的GROUP BY和聚合函数。")
    
    constraints_text = "\n".join(constraints)
    
    prompt = f"""
根据以下数据库架构:
{tables_info}

将这个问题转换为SQL查询:
"{question}"

注意事项:
1. 对于ENUM类型的字段，必须使用上述指定的枚举值之一
2. ENUM值必须用双引号括起来
3. 不要使用不必要的表别名，除非查询涉及同一表的多次引用
4. 不要为列添加不必要的别名，特别是聚合函数结果
5. 使用双引号而非单引号表示字符串
6. 不要使用AS关键字，除非绝对必要
7. 保持SQL简洁，不要添加不必要的空格或格式化

{constraints_text}

请生成适用于MySQL数据库的SQL查询。
"""
    # 使用知识库增强提示词
    return knowledge_base.enhance_prompt(prompt, question)
# 仅返回SQL查询语句，不要包含任何解释或额外文本
def call_tongyi_api(prompt):
    """调用通义千问API将自然语言转换为SQL"""
    try:
        response = Generation.call(
            model='qwen-max',
            prompt=prompt,
            api_key=config.TONGYI_API_KEY
        )
        
        if response.status_code == 200:
            return response.output.text
        else:
            print(f"调用通义千问API错误: {response.code}, {response.message}")
            return None
    except Exception as e:
        print(f"调用通义千问API错误: {e}")
        return None

def call_deepseek_api(prompt):
    """调用DeepSeek API将自然语言转换为SQL"""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {config.DEEPSEEK_API_KEY}"
    }
    
    data = {
        "model": "deepseek-chat",
        "messages": [
            {"role": "system", "content": config.PROMPT_TEMPLATE_IDENTITY},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.3,
        "max_tokens": 500
    }
    
    try:
        response = requests.post(
            "https://api.deepseek.com/v1/chat/completions", 
            headers=headers, 
            json=data
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]
    except Exception as e:
        print(f"调用DeepSeek API错误: {e}")
        return None

def generate_sql(question, db_schema, options):
    """使用选定的AI模型将自然语言转换为SQL"""
    prompt = generate_prompt(question, db_schema, options)
    print(f"Generated prompt: {prompt}")  # 调试输出
    
    response = None
    if config.ACTIVE_AI_MODEL == "deepseek":
        response = call_deepseek_api(prompt)
    else:
        response = call_tongyi_api(prompt)
    print(f"AI response: {response}")  # 调试输出
    
    return str(response) if response is not None else None

def extract_sql_from_ai_response(response):
    """从AI模型响应中提取SQL查询"""
    if not response:
        return "AI模型无相应，请检查原因......"  # 没有响应
    
    # 确保response是字符串
    if not isinstance(response, str):
        response = str(response)
    
    # 尝试从SQL代码块中提取
    if "```sql" in response:
        parts = response.split("```sql", 1)
        if len(parts) > 1:
            sql_part = parts[1].split("```", 1)
            if len(sql_part) > 0:
                return sql_part[0].strip()
    
    # 尝试从一般代码块中提取
    if "```" in response:
        parts = response.split("```", 1)
        if len(parts) > 1:
            sql_part = parts[1].split("```", 1)
            if len(sql_part) > 0:
                return sql_part[0].strip()
    
    # 尝试从反引号中提取
    if "`" in response:
        parts = response.split("`")
        if len(parts) > 2:
            return parts[1].strip()
    
    # 如果没有格式化，返回整个响应
    return response.strip()

def generate_explanation(question, results, sql, error=None):
    """生成查询结果的解释"""
    if error:
        return f"查询执行失败: {error}"
    
    result_count = len(results)
    
    if result_count == 0:
        return f"未找到与'{question}'相关的结果。"
    else:
        return f"找到{result_count}条与'{question}'相关的结果。"

def refine_sql_with_feedback(question, initial_sql, error=None, results=None, db_schema=None):
    """
    基于第一次执行的结果或错误，优化SQL查询
    
    Args:
        question (str): 原始用户问题
        initial_sql (str): 第一次生成的SQL
        error (str, optional): 执行错误信息
        results (list, optional): 第一次执行的结果
        db_schema (DatabaseSchema): 数据库模式
    
    Returns:
        str: 优化后的SQL查询
    """
    # 构建反馈提示词
    feedback_prompt = f"""
根据用户问题:
"{question}"

我生成了以下SQL查询:
```sql
{initial_sql}
```
"""
    
    # 添加错误信息或结果反馈
    if error:
        # 提供更详细的错误分析和修复建议
        feedback_prompt += f"""
但执行时遇到以下错误:
{error}

请仔细分析错误原因，可能是:
1. 语法错误: 检查SQL语法，特别是关键字、括号匹配、引号使用
2. 表名错误: 确保表名与数据库模式匹配，注意大小写
3. 列名错误: 确保列名存在于相应的表中，注意大小写
4. 数据类型不匹配: 确保条件中的值类型与列类型匹配
5. JOIN条件错误: 检查JOIN条件是否正确连接相关表

请修正SQL查询中的错误，确保语法正确并与数据库模式匹配。
"""
    elif results is not None and len(results) == 0:
        # 提供更详细的空结果分析和优化建议
        feedback_prompt += """
查询执行成功，但没有返回任何结果。这可能是因为:
1. 查询条件过于严格: 尝试放宽WHERE条件，移除一些限制性条件
2. 表名或列名可能有误: 仔细检查所有表名和列名是否正确
3. 查询逻辑可能与用户意图不符: 重新分析用户问题，确保SQL查询正确反映用户意图
4. JOIN条件可能过滤掉了所有结果: 检查JOIN条件是否正确
5. 数据库中可能没有满足条件的数据: 考虑数据的实际情况，调整查询条件

请重新分析用户问题，生成一个能返回相关结果的SQL查询。尝试以下策略:
- 简化查询条件
- 检查表名和列名
- 重新思考查询逻辑
- 考虑使用不同的JOIN类型(如LEFT JOIN)
- 确保聚合函数使用正确
"""
    
    # 添加数据库模式信息，并提供更详细的模式解释
    if db_schema:
        tables_info = ""
        # 添加表之间的关系信息
        relationships = []
        
        for table in db_schema.tables:
            columns_info = []
            primary_keys = []
            foreign_keys = []
            
            for col in table.columns:
                # 收集主键信息
                if hasattr(col, 'is_primary_key') and col.is_primary_key:
                    primary_keys.append(col.name)
                
                # 收集外键信息
                if hasattr(col, 'references') and col.references:
                    ref_table, ref_col = col.references.split('.')
                    foreign_keys.append(f"{col.name} -> {ref_table}.{ref_col}")
                    relationships.append(f"{table.name}.{col.name} 引用 {ref_table}.{ref_col}")
                
                # 格式化列信息
                if hasattr(col, 'enum_values') and col.enum_values:
                    enum_values_str = ", ".join(f"'{v}'" for v in col.enum_values.keys())
                    columns_info.append(f"{col.name} ({col.type}): {col.description} [可选值: {enum_values_str}]")
                else:
                    col_type = col.type if hasattr(col, 'type') else "未知类型"
                    col_desc = col.description if hasattr(col, 'description') else ""
                    columns_info.append(f"{col.name} ({col_type}){': ' + col_desc if col_desc else ''}")
            
            # 添加主键和外键信息
            pk_info = f"主键: {', '.join(primary_keys)}" if primary_keys else ""
            fk_info = f"外键: {', '.join(foreign_keys)}" if foreign_keys else ""
            
            tables_info += f"表: {table.name}\n"
            tables_info += f"列: {', '.join(columns_info)}\n"
            if pk_info:
                tables_info += f"{pk_info}\n"
            if fk_info:
                tables_info += f"{fk_info}\n"
            tables_info += "\n"
        
        # 添加表关系信息
        if relationships:
            tables_info += "表关系:\n"
            tables_info += "\n".join(relationships)
            tables_info += "\n\n"
        
        feedback_prompt += f"""
数据库架构:
{tables_info}

请特别注意表之间的关系，确保JOIN条件正确。
"""
    
    # 添加查询分析和优化建议
    feedback_prompt += """
请分析原始SQL查询中可能存在的问题:

1. 语法问题:
   - 检查关键字拼写
   - 确保括号正确匹配
   - 检查引号使用是否一致

2. 结构问题:
   - 确保SELECT子句中的列在FROM子句的表中存在
   - 检查GROUP BY子句是否包含所有非聚合列
   - 确保HAVING子句只用于过滤聚合结果

3. 语义问题:
   - 重新分析用户问题，确保SQL查询正确反映用户意图
   - 检查WHERE条件是否正确表达了用户的筛选需求
   - 确保JOIN条件正确连接相关表

4. 优化策略:
   - 简化复杂条件
   - 确保使用正确的JOIN类型
   - 考虑使用子查询或WITH子句分解复杂查询

现在，请生成一个优化后的SQL查询，确保它能正确执行并返回相关结果。
只返回SQL语句，不要包含任何解释或额外文本。
"""
    
    # 使用选定的AI模型生成优化后的SQL
    if config.ACTIVE_AI_MODEL == "deepseek":
        response = call_deepseek_api(feedback_prompt)
    else:
        response = call_tongyi_api(feedback_prompt)
    
    # 如果没有得到有效响应，返回原始SQL
    if not response:
        return initial_sql
        
    return response
